import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
from tensorflow.contrib.factorization import KMeans

mnist = input_data.read_data_sets("input_data/", one_hot=True)
full_data_x = mnist.train.images

# 定义超参数
num_steps = 500
batch_size = 1024
k = 25
num_features = 784
num_classes = 10

# 定义输入输出
X = tf.placeholder(tf.float32, shape=[None, num_features])
Y = tf.placeholder(tf.float32, shape=[None, num_classes])
# 使用tensorflow集成的方法
kmeans = KMeans(inputs=X, num_clusters=k, distance_metric='cosine', use_mini_batch=True)
# 构建kmeans图
training_graph = kmeans.training_graph()

# tensorflow版本不同所以启用不同参数 tensorflow1.4+多了一个参数
if len(training_graph) > 6: # Tensorflow 1.4+
    (all_scores, cluster_idx, scores, cluster_centers_initialized,
     cluster_centers_var, init_op, train_op) = training_graph
else:
    (all_scores, cluster_idx, scores, cluster_centers_initialized,
     init_op, train_op) = training_graph

# 将族固定成一个元组
cluster_idx = cluster_idx[0] # fix for cluster_idx being a tuple
# 计算平均距离
avg_distance = tf.reduce_mean(scores)

# 初始化
init_var = tf.global_variables_initializer()

sess = tf.Session()
sess.run(init_var, feed_dict={X:full_data_x})
sess.run(init_op, feed_dict={X:full_data_x})

# 开始训练
for step in range(1, num_steps+1):

    _, d, idx = sess.run([train_op, avg_distance, cluster_idx], feed_dict={X:full_data_x})
    print("step %i,AvgDistance is %f" % (step, d))

# 载入分类
counts = np.zeros(shape=(k, num_classes))
for i in range(len(idx)):
    counts[idx[i]] += mnist.train.labels[i]

# 指派给频繁出现标签
labels_map = [np.argmax(c) for c in counts]
labels_map = tf.convert_to_tensor(labels_map)

# 按照组内顺序返回cluster_id行
cluster_label = tf.nn.embedding_lookup(labels_map, cluster_idx)

# 计算准确率
correct_prediction = tf.equal(cluster_label, tf.cast(tf.argmax(Y, 1), tf.int32))
accuracy_op = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

# 测试模型
test_x, test_y = mnist.test.images, mnist.test.labels
print("Test Accuracy:", sess.run(accuracy_op, feed_dict={X: test_x, Y: test_y}))




